package xyz.lwm.lazycat.aop.support;

import net.bytebuddy.ByteBuddy;
import net.bytebuddy.description.ByteCodeElement;
import net.bytebuddy.description.modifier.Visibility;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.implementation.FieldAccessor;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatcher;
import xyz.lwm.lazycat.aop.AopUtil;
import xyz.lwm.lazycat.aop.Aspect;
import xyz.lwm.lazycat.aop.LazycatProxy;
import xyz.lwm.lazycat.aop.info.ClassInfo;
import xyz.lwm.lazycat.aop.info.MethodInfo;
import xyz.lwm.lazycat.aop.InvocationHandler;
import xyz.lwm.lazycat.utility.CollectionUtil;
import xyz.lwm.lazycat.utility.Pair;
import xyz.lwm.lazycat.utility.reflect.ReflectUtil;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

import static net.bytebuddy.matcher.ElementMatchers.is;
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
import static net.bytebuddy.matcher.ElementMatchers.not;

/**
 * Use byte buddy to create proxy object
 *
 * @author lwm
 */
public abstract class ProxyUtil {

    private static final String DELEGATE_FIELD_NAME = "$lazycat_delegate$";

    private static final String DELEGATE_CLASS_NAME_SUFFIX = "$$LazycatProxy";

    private static final AtomicInteger INDEX = new AtomicInteger(100);

    private static final Supplier<String> JDK_PROXY_CLASS_NAME = () -> "LazycatJdkProxy" + INDEX.incrementAndGet() + DELEGATE_CLASS_NAME_SUFFIX;

    private static final Object LOCK = new Object();

    /**
     * create proxy object for target object
     *
     * @param target target object
     * @param aspects aspects collection
     * @param <T> generic type
     * @return proxy object
     */
    public static <T> T createProxy(T target, Collection<Aspect> aspects) {
        return createProxy0(target, null, aspects);
    }

    /**
     * Handle the proxy object recursively
     */
    @SuppressWarnings("unchecked")
    private static <T> T createProxy0(T target, T unwrapped, Collection<Aspect> aspects) {
        unwrapped = unwrapped == null ? target : unwrapped;
        if (AopUtil.isJdkProxy(unwrapped)) {
            Class<?>[] interfaces = unwrapped.getClass().getInterfaces();
            return createLazycatProxyForJdkTarget(target, interfaces, aspects);

        } else if (AopUtil.isAopProxy(target)) {
            return createProxy0(target, AopUtil.unwrap(target), aspects);

        } else {
            return createLazycatProxy(target, (Class<? super T>) target.getClass(), aspects);
        }
    }

    /**
     * create proxy object for original object or lazycat proxy object
     */
    @SuppressWarnings("unchecked")
    private static <T> T createLazycatProxy(T target, Class<? super T> clz, Collection<Aspect> aspects) {
        // get methods for aspects
        List<Pair<Method, List<InvocationHandler>>> methods = getMethods(Collections.singleton(clz), aspects);
        if (CollectionUtil.isEmpty(methods)) {
            return target;
        }
        // make and load proxy class
        DynamicType.Builder.MethodDefinition.ReceiverTypeDefinition<?> definition =
                new ByteBuddy().subclass(clz).implement(LazycatProxy.class)
                        .name(clz.getName() + DELEGATE_CLASS_NAME_SUFFIX)
                        .defineField(DELEGATE_FIELD_NAME, clz, Visibility.PRIVATE)
                        .method(not(isDeclaredBy(Object.class)))
                        .intercept(MethodDelegation.toField(DELEGATE_FIELD_NAME))
                        .method(isDeclaredBy(LazycatProxy.class))
                        .intercept(FieldAccessor.ofField(DELEGATE_FIELD_NAME));
        // add interceptor for proxied methods
        for (Pair<Method, List<InvocationHandler>> method : methods) {
            Interceptor interceptor = new Interceptor(target, method.getRight().toArray(new InvocationHandler[0]));
            definition = definition.method(is(method.getLeft())).intercept(MethodDelegation.to(interceptor));
        }
        Class<?> proxyClz = definition.make().load(clz.getClassLoader()).getLoaded();

        // create proxy object
        T proxy = (T) ReflectUtil.newInstance(proxyClz);
        // set delegate field by reflect
        ReflectUtil.setValue(proxy, DELEGATE_FIELD_NAME, target);

        return proxy;
    }

    /**
     * create proxy object for jdk proxy object
     */
    @SuppressWarnings("unchecked")
    private static <T> T createLazycatProxyForJdkTarget(T target, Class<?>[] interfaces, Collection<Aspect> aspects) {
        // get methods for aspects
        List<Pair<Method, List<InvocationHandler>>> methods = getMethods(Arrays.asList(interfaces), aspects);
        if (CollectionUtil.isEmpty(methods)) {
            return target;
        }

        // get method matcher
        ElementMatcher.Junction<ByteCodeElement> matcher = isDeclaredBy(interfaces[0]);
        for (int i = 1; i < interfaces.length; i++) {
            matcher = matcher.or(isDeclaredBy(interfaces[i]));
        }

        // make and load proxy class
        DynamicType.Builder.MethodDefinition.ReceiverTypeDefinition<?> definition =
                new ByteBuddy().subclass(Object.class)
                        .name(JDK_PROXY_CLASS_NAME.get())
                        .implement(LazycatProxy.class)
                        .implement(interfaces)
                        .defineField(DELEGATE_FIELD_NAME, target.getClass(), Visibility.PRIVATE)
                        .method(isDeclaredBy(LazycatProxy.class))
                        .intercept(FieldAccessor.ofField(DELEGATE_FIELD_NAME))
                        .method(matcher)
                        .intercept(MethodDelegation.toField(DELEGATE_FIELD_NAME));
        // add interceptor for proxied methods
        for (Pair<Method, List<InvocationHandler>> method : methods) {
            Interceptor interceptor = new Interceptor(target, method.getRight().toArray(new InvocationHandler[0]));
            definition = definition.method(is(method.getLeft())).intercept(MethodDelegation.to(interceptor));
        }
        Class<?> proxyClz = definition.make().load(target.getClass().getClassLoader()).getLoaded();

        // create proxy object
        T proxy = (T) ReflectUtil.newInstance(proxyClz);
        // set delegate field by reflect
        ReflectUtil.setValue(proxy, DELEGATE_FIELD_NAME, target);

        return proxy;
    }

    private static List<Pair<Method, List<InvocationHandler>>> getMethods(Collection<Class<?>> classes, Collection<Aspect> aspects) {
        Map<String, Pair<Method, List<InvocationHandler>>> methodMap = new LinkedHashMap<>();
        for (Class<?> clz : classes) {
            ClassInfo classInfo = ClassInfo.of(clz);
            for (MethodInfo methodInfo : classInfo) {
                String key = methodInfo.getMethod().toString();
                if (!methodMap.containsKey(key)) {
                    List<InvocationHandler> list = new ArrayList<>();
                    for (Aspect aspect : aspects) {
                        synchronized (LOCK) {
                            if (aspect.matcher().matches(methodInfo)) {
                                list.add(aspect.interceptor());
                            }
                        }
                    }
                    if (!list.isEmpty()) {
                        methodMap.put(key, Pair.of(methodInfo.getMethod(), list));
                    }
                }
            }
        }
        return CollectionUtil.collect(methodMap.values(), ArrayList::new);
    }

}
